from dotenv import load_dotenv
import pymysql
import pyodbc
import re
from pymysql.err import OperationalError as MySQLOperationalError
from pyodbc import OperationalError as SQLServerOperationalError
from typing import List, Optional, Dict, Any

load_dotenv()

class DatabaseConnectionError(Exception):
    """数据库连接错误的自定义异常"""
    pass

def test_mysql_connection(host, port, user, password, database, extra_params=None):
    """
    测试MySQL数据库连接
    
    如果连接成功返回True，如果失败则抛出DatabaseConnectionError异常
    """
    try:
        connection = pymysql.connect(
            host=host,
            port=int(port),
            user=user,
            password=password,
            database=database,
            connect_timeout=30
        )
        with connection.cursor():
            pass  # 仅测试连接
        connection.close()
        return True
    except MySQLOperationalError as e:
        raise DatabaseConnectionError(f"MySQL连接失败: {str(e)}")
    except Exception as e:
        raise DatabaseConnectionError(f"连接MySQL时发生意外错误: {str(e)}")

def test_sqlserver_connection(host, port, user, password, database, extra_params=None):
    """
    测试SQL Server数据库连接
    
    如果连接成功返回True，如果失败则抛出DatabaseConnectionError异常
    """
    try:
        # 构建连接字符串
        conn_str = f"DRIVER={{ODBC Driver 18 for SQL Server}};SERVER={host},{port};DATABASE={database};UID={user};PWD={password};TrustServerCertificate=yes;"
        
        # 添加额外的连接参数
        if extra_params:
            conn_str += f";{extra_params}"
            
        connection = pyodbc.connect(conn_str)
        with connection.cursor():
            pass  # 仅测试连接
        connection.close()
        return True
    except SQLServerOperationalError as e:
        raise DatabaseConnectionError(f"SQL Server连接失败: {str(e)}")
    except Exception as e:
        raise DatabaseConnectionError(f"连接SQL Server时发生意外错误: {str(e)}")

def test_fabric_sqlserver_connection(host, port, user, password, database, extra_params=None):
    """
    测试Microsoft Fabric Data Warehouse数据库连接
    
    如果连接成功返回True，如果失败则抛出DatabaseConnectionError异常
    
    注意：当使用ActiveDirectoryAccessToken认证时，password参数应该包含访问令牌
    """
    try:
        # Microsoft Fabric Data Warehouse使用特殊连接字符串
        conn_str = (
            f"DRIVER={{ODBC Driver 18 for SQL Server}};"  # 尽可能使用18版驱动
            f"SERVER={host};"  # Fabric端点通常不需要端口
            f"DATABASE={database};"
            f"Uid={user};"
            f"Pwd={password};"
            f"Encrypt=yes;"
            f"TrustServerCertificate=yes;"
            f"Authentication=ActiveDirectoryServicePrincipal;"
            f"ssl=true;"
            f"Connection Timeout=300;"
        )
        # 添加额外的连接参数
        if extra_params:
            conn_str += f"{extra_params}"
        connection = pyodbc.connect(conn_str)
        with connection.cursor():
            pass  # 仅测试连接 
        connection.close()
        return True
    except SQLServerOperationalError as e:
        raise DatabaseConnectionError(f"Fabric SQL Server连接失败: {str(e)}")
    except Exception as e:
        raise DatabaseConnectionError(f"连接Microsoft Fabric Data Warehouse时发生意外错误: {str(e)}")

def test_database_connection(db_type, host, port, user, password, database, extra_params=None):
    """
    根据数据库类型测试数据库连接
    """
    if db_type.lower() == "mysql":
        return test_mysql_connection(host, port, user, password, database, extra_params)
    elif db_type.lower() == "sqlserver":
        return test_sqlserver_connection(host, port, user, password, database, extra_params)
    elif db_type.lower() == "fabric_sqlserver":
        return test_fabric_sqlserver_connection(host, port, user, password, database, extra_params)
    else:
        raise DatabaseConnectionError(f"不支持的数据库类型: {db_type}")

def get_mysql_tables(host, port, user, password, database, extra_params=None) -> List[str]:
    """
    获取MySQL数据库中的表列表
    
    Args:
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        extra_params: 额外的连接参数
    
    Returns:
        表名列表
    """
    try:
        connection = pymysql.connect(
            host=host,
            port=int(port),
            user=user,
            password=password,
            database=database,
            connect_timeout=30
        )
        
        tables = []
        with connection.cursor() as cursor:
            cursor.execute("SHOW TABLES")
            for row in cursor.fetchall():
                tables.append(row[0])
                
        connection.close()
        return tables
    except Exception as e:
        raise DatabaseConnectionError(f"获取MySQL表列表失败: {str(e)}")

def get_sqlserver_tables(host, port, user, password, database, extra_params=None) -> List[str]:
    """
    获取SQL Server数据库中的表列表
    
    Args:
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        extra_params: 额外的连接参数
    
    Returns:
        表名列表
    """
    try:
        # 构建连接字符串
        conn_str = f"DRIVER={{ODBC Driver 18 for SQL Server}};SERVER={host},{port};DATABASE={database};UID={user};PWD={password};TrustServerCertificate=yes;"
        
        # 添加额外的连接参数
        if extra_params:
            conn_str += f";{extra_params}"
        
        connection = pyodbc.connect(conn_str)
        
        tables = []
        with connection.cursor() as cursor:
            # 查询用户表
            query = """
                SELECT TABLE_NAME 
                FROM INFORMATION_SCHEMA.TABLES 
                WHERE TABLE_TYPE = 'BASE TABLE' 
                ORDER BY TABLE_NAME
            """
            cursor.execute(query)
            for row in cursor.fetchall():
                tables.append(row[0])
                
        connection.close()
        return tables
    except Exception as e:
        raise DatabaseConnectionError(f"获取SQL Server表列表失败: {str(e)}")

def get_fabric_sqlserver_tables(host, port, user, password, database, extra_params=None) -> List[str]:
    """
    获取Microsoft Fabric Data Warehouse数据库中的表列表
    
    Args:
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        extra_params: 额外的连接参数
    
    Returns:
        表名列表
    """
    try:
        # Microsoft Fabric Data Warehouse使用特殊连接字符串
        conn_str = (
            f"DRIVER={{ODBC Driver 18 for SQL Server}};"
            f"SERVER={host};"
            f"DATABASE={database};"
            f"Uid={user};"
            f"Pwd={password};"
            f"Encrypt=yes;"
            f"TrustServerCertificate=yes;"
            f"Authentication=ActiveDirectoryServicePrincipal;"
            f"ssl=true;"
            f"Connection Timeout=300;"
        )
        
        # 添加额外的连接参数
        if extra_params:
            conn_str += f"{extra_params}"
            
        connection = pyodbc.connect(conn_str)
        
        tables = []
        with connection.cursor() as cursor:
            # 查询用户表
            query = """
                SELECT TABLE_NAME 
                FROM INFORMATION_SCHEMA.TABLES 
                WHERE TABLE_TYPE = 'BASE TABLE' 
                ORDER BY TABLE_NAME
            """
            cursor.execute(query)
            for row in cursor.fetchall():
                tables.append(row[0])
                
        connection.close()
        return tables
    except Exception as e:
        raise DatabaseConnectionError(f"获取Microsoft Fabric Data Warehouse表列表失败: {str(e)}")

def get_database_tables(db_type, host, port, user, password, database, extra_params=None) -> List[str]:
    """
    根据数据库类型获取表列表
    
    Args:
        db_type: 数据库类型
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        extra_params: 额外的连接参数
    
    Returns:
        表名列表
    """
    if db_type.lower() == "mysql":
        return get_mysql_tables(host, port, user, password, database, extra_params)
    elif db_type.lower() == "sqlserver":
        return get_sqlserver_tables(host, port, user, password, database, extra_params)
    elif db_type.lower() == "fabric_sqlserver":
        return get_fabric_sqlserver_tables(host, port, user, password, database, extra_params)
    else:
        raise DatabaseConnectionError(f"不支持的数据库类型: {db_type}")

def get_mysql_columns(host, port, user, password, database, table_name, extra_params=None) -> List[str]:
    """
    获取MySQL表中的字段列表
    
    Args:
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        table_name: 表名
        extra_params: 额外的连接参数
    
    Returns:
        字段名列表
    """
    try:
        connection = pymysql.connect(
            host=host,
            port=int(port),
            user=user,
            password=password,
            database=database,
            connect_timeout=30
        )
        
        columns = []
        with connection.cursor() as cursor:
            cursor.execute(f"SHOW COLUMNS FROM `{table_name}`")
            for row in cursor.fetchall():
                columns.append(row[0])
                
        connection.close()
        return columns
    except Exception as e:
        raise DatabaseConnectionError(f"获取MySQL表字段列表失败: {str(e)}")

def get_sqlserver_columns(host, port, user, password, database, table_name, extra_params=None) -> List[str]:
    """
    获取SQL Server表中的字段列表
    
    Args:
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        table_name: 表名
        extra_params: 额外的连接参数
    
    Returns:
        字段名列表
    """
    try:
        # 构建连接字符串
        conn_str = f"DRIVER={{ODBC Driver 18 for SQL Server}};SERVER={host},{port};DATABASE={database};UID={user};PWD={password};TrustServerCertificate=yes;"
        
        # 添加额外的连接参数
        if extra_params:
            conn_str += f";{extra_params}"
            
        connection = pyodbc.connect(conn_str)
        
        columns = []
        with connection.cursor() as cursor:
            query = f"""
                SELECT COLUMN_NAME
                FROM INFORMATION_SCHEMA.COLUMNS
                WHERE TABLE_NAME = '{table_name}'
                ORDER BY ORDINAL_POSITION
            """
            cursor.execute(query)
            for row in cursor.fetchall():
                columns.append(row[0])
                
        connection.close()
        return columns
    except Exception as e:
        raise DatabaseConnectionError(f"获取SQL Server表字段列表失败: {str(e)}")

def get_fabric_sqlserver_columns(host, port, user, password, database, table_name, extra_params=None) -> List[str]:
    """
    获取Microsoft Fabric Data Warehouse表中的字段列表
    
    Args:
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        table_name: 表名
        extra_params: 额外的连接参数
    
    Returns:
        字段名列表
    """
    try:
        # Microsoft Fabric Data Warehouse使用特殊连接字符串
        conn_str = (
            f"DRIVER={{ODBC Driver 18 for SQL Server}};"
            f"SERVER={host};"
            f"DATABASE={database};"
            f"Uid={user};"
            f"Pwd={password};"
            f"Encrypt=yes;"
            f"TrustServerCertificate=yes;"
            f"Authentication=ActiveDirectoryServicePrincipal;"
            f"ssl=true;"
            f"Connection Timeout=300;"
        )
        
        # 添加额外的连接参数
        if extra_params:
            conn_str += f"{extra_params}"
            
        connection = pyodbc.connect(conn_str)
        
        columns = []
        with connection.cursor() as cursor:
            query = f"""
                SELECT COLUMN_NAME
                FROM INFORMATION_SCHEMA.COLUMNS
                WHERE TABLE_NAME = '{table_name}'
                ORDER BY ORDINAL_POSITION
            """
            cursor.execute(query)
            for row in cursor.fetchall():
                columns.append(row[0])
                
        connection.close()
        return columns
    except Exception as e:
        raise DatabaseConnectionError(f"获取Microsoft Fabric Data Warehouse表字段列表失败: {str(e)}")

def get_table_columns(db_type, host, port, user, password, database, table_name, extra_params=None) -> List[str]:
    """
    根据数据库类型获取表字段列表
    
    Args:
        db_type: 数据库类型
        host: 数据库主机名
        port: 端口号
        user: 用户名
        password: 密码
        database: 数据库名
        table_name: 表名
        extra_params: 额外的连接参数
    
    Returns:
        字段名列表
    """
    if db_type.lower() == "mysql":
        return get_mysql_columns(host, port, user, password, database, table_name, extra_params)
    elif db_type.lower() == "sqlserver":
        return get_sqlserver_columns(host, port, user, password, database, table_name, extra_params)
    elif db_type.lower() == "fabric_sqlserver":
        return get_fabric_sqlserver_columns(host, port, user, password, database, table_name, extra_params)
    else:
        raise DatabaseConnectionError(f"不支持的数据库类型: {db_type}")